#include "FireworkParticle.h"

#include <iostream>

#include <vsg/all.h>

#include "utilities.h"
#include "FireworkParticle.h"
#include "DataSingleton.h"

void FireworkParticle::update(const std::chrono::nanoseconds& delt) {
	if (isDone()) {
		return;
	}
	if (m_trailLength > 0) {
		m_oldPos.push_back(m_pos);
	}
	if (m_oldPos.size() > m_trailLength) {
		m_oldPos.pop_front();
	}

	Particle::update(delt);

	if (isDone()) {
		invalid();
		return;
	}

	for (auto& color : *m_trailColors) {
		color.set(m_color.r, m_color.g, m_color.b, m_color.a);
	}
	m_trailColors->dirty();

	if (m_oldPos.size() > 0) {
		vsg::dvec3 axisZ = *m_oldPos.cbegin() - m_pos;
		double lengthZ = length(axisZ);
		auto normalZ = normalize(axisZ);
		auto a = cross(normalZ, { 1,0,0 });
		if (a.x == 0 && a.y == 0 && a.z == 0) {
			a = cross(a, { 0,1,0 });
		}
		auto b = cross(a, normalZ);
		a = normalize(a);
		b = normalize(b);
		vsg::dmat4 modelToWorld({ a, 0 }, { b, 0 }, { normalZ,0 }, { 0,0,0,1 });
		m_trailMatrixTransform->matrix = vsg::translate(m_pos) * modelToWorld * vsg::scale(m_radius, m_radius, lengthZ);
	}

}

vsg::ref_ptr<vsg::Node> FireworkParticle::getNode()
{
	m_switch = vsg::Switch::create();
	auto particleScene = Particle::getNode();
	m_trailMatrixTransform = vsg::MatrixTransform::create();
	m_switch->addChild(false, particleScene);

	using namespace vsg;
	vsg::Paths searchPaths = vsg::getEnvPaths("VSG_FILE_PATH");
	vsg::ref_ptr<vsg::ShaderStage> vertexShader = vsg::ShaderStage::read(VK_SHADER_STAGE_VERTEX_BIT, "main", vsg::findFile("shaders/firework_vsg_vert.spv", searchPaths));
	vsg::ref_ptr<vsg::ShaderStage> fragmentShader = vsg::ShaderStage::read(VK_SHADER_STAGE_FRAGMENT_BIT, "main", vsg::findFile("shaders/firework_vsg_frag.spv", searchPaths));
	if (!vertexShader || !fragmentShader)
	{
		std::cout << "Could not create shaders." << std::endl;
	}

	vsg::PushConstantRanges pushConstantRanges{
		{VK_SHADER_STAGE_VERTEX_BIT, 0, 128} // projection view, and model matrices, actual push constant calls automatically provided by the VSG's DispatchTraversal
	};

	vsg::VertexInputState::Bindings vertexBindingsDescriptions{
		VkVertexInputBindingDescription{0, sizeof(vsg::vec3), VK_VERTEX_INPUT_RATE_VERTEX}, // vertex data
		VkVertexInputBindingDescription{1, sizeof(vsg::vec3), VK_VERTEX_INPUT_RATE_VERTEX}, // colour data
	};

	vsg::VertexInputState::Attributes vertexAttributeDescriptions{
		VkVertexInputAttributeDescription{0, 0, VK_FORMAT_R32G32B32_SFLOAT, 0}, // vertex data
		VkVertexInputAttributeDescription{1, 1, VK_FORMAT_R32G32B32_SFLOAT, 0}, // colour data
	};

	vsg::GraphicsPipelineStates pipelineStates{
		vsg::VertexInputState::create(vertexBindingsDescriptions, vertexAttributeDescriptions),
		vsg::InputAssemblyState::create(),
		vsg::RasterizationState::create(),
		vsg::MultisampleState::create(),
		vsg::ColorBlendState::create(),
		vsg::DepthStencilState::create() };

	auto pipelineLayout = vsg::PipelineLayout::create(vsg::DescriptorSetLayouts{}, pushConstantRanges);
	auto graphicsPipeline = vsg::GraphicsPipeline::create(pipelineLayout, vsg::ShaderStages{ vertexShader, fragmentShader }, pipelineStates);
	auto bindGraphicsPipeline = vsg::BindGraphicsPipeline::create(graphicsPipeline);

	// create StateGroup as the root of the scene/command graph to hold the GraphicsProgram, and binding of Descriptors to decorate the whole graph
	auto scenegraph = vsg::StateGroup::create();
	m_switch->addChild(false, scenegraph);
	scenegraph->add(bindGraphicsPipeline);


	// add transform to root of the scene graph
	scenegraph->addChild(m_trailMatrixTransform);

	vec3 dx = { 0.5f, 0.0f, 0.0f };
	vec3 dy = { 0.0f, 0.5f, 0.0f };
	vec3 dz = { 0.0f, 0.0f, 0.5f };

	vec3 bottom = { 0.f, 0.f, 0.f };
	vec3 top = { 0.f, 0.f, 1.f };

	ref_ptr<vec3Array> vertices;
	ref_ptr<ushortArray> indices;

	auto edge = [&](float alpha) -> vec3 {
		return dy * (cosf(alpha)) - dx * (sinf(alpha));
	};

	auto normal = [&](float alpha) -> vec3 {
		float delta = 0.001f;
		vec3 before = edge(alpha - delta);
		vec3 mid = edge(alpha);
		vec3 after = edge(alpha + delta);
		return normalize(cross(after - before, dz - mid));
	};

	float alpha = 0.0f;
	vec3 v = edge(alpha);
	vec3 n = normal(alpha);

	bool withEnds = true;

	unsigned int num_columns = 20;
	unsigned int num_vertices = num_columns * 2;
	unsigned int num_indices = (num_columns - 1) * 3;

	if (withEnds)
	{
		num_vertices += num_columns;
		num_indices += (num_columns - 2) * 3;
	}

	vertices = vec3Array::create(num_vertices);
	indices = ushortArray::create(num_indices);

	m_trailColors = vsg::vec4Array::create(num_vertices, m_color);
	m_trailColors->properties.dataVariance = DYNAMIC_DATA_TRANSFER_AFTER_RECORD;

	vertices->set(0, bottom + v);
	vertices->set(num_columns * 2 - 2, bottom + v);

	vertices->set(1, top);
	vertices->set(num_columns * 2 - 1, top);

	for (unsigned int c = 1; c < num_columns - 1; ++c)
	{
		unsigned int vi = c * 2;
		float r = float(c) / float(num_columns - 1);
		alpha = (r) * 2.0f * PIf;
		v = edge(alpha);
		n = normal(alpha);

		vertices->set(vi, bottom + v);
		vertices->set(vi + 1, top);
	}

	unsigned int i = 0;
	for (unsigned int c = 0; c < num_columns - 1; ++c)
	{
		unsigned lower = c * 2;
		unsigned upper = lower + 1;

		indices->set(i++, lower);
		indices->set(i++, lower + 2);
		indices->set(i++, upper);
	}

	if (withEnds)
	{
		unsigned int bottom_i = num_columns * 2;
		v = edge(0.0f);
		vec3 bottom_n = normalize(-dz);

		vertices->set(bottom_i, bottom + v);
		vertices->set(bottom_i + num_columns - 1, bottom + v);

		for (unsigned int c = 1; c < num_columns - 1; ++c)
		{
			float r = float(c) / float(num_columns - 1);
			alpha = (r) * 2.0f * PIf;
			v = edge(alpha);

			unsigned int vi = bottom_i + c;
			vertices->set(vi, bottom + v);
		}

		for (unsigned int c = 0; c < num_columns - 2; ++c)
		{
			indices->set(i++, bottom_i + c);
			indices->set(i++, bottom_i + num_columns - 1);
			indices->set(i++, bottom_i + c + 1);
		}
	}

	// setup geometry
	auto drawCommands = vsg::Commands::create();
	drawCommands->addChild(vsg::BindVertexBuffers::create(0, vsg::DataList{ vertices, m_trailColors }));
	drawCommands->addChild(vsg::BindIndexBuffer::create(indices));
	drawCommands->addChild(vsg::DrawIndexed::create(num_indices, 1, 0, 0, 0));

	m_trailMatrixTransform->addChild(drawCommands);
	return m_switch;
}

void FireworkParticle::activate()
{
	Particle::activate();
	m_switch->setAllChildren(true);
	m_oldPos.clear();
}

void FireworkParticle::invalid()
{
	Particle::invalid();
	m_switch->setAllChildren(false);
	m_oldPos.clear();
}

void FireworkParticle::initNode()
{
}
